package com.archgeek.java.udp;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;

/**
 * @author pizhihui
 * @date 2023-08-07 21:13
 */
public class JobManagerDiscoverService {


    public static void main(String[] args) throws IOException {

        DatagramSocket serverSocket = new DatagramSocket(6123, InetAddress.getLocalHost());
        System.out.println("server socket bind to " + serverSocket.getLocalSocketAddress());

        int ipcPort = 6124;
        int magicNumber = 0;

        new Thread(() -> {
            try {
                while (true) {

                    final DatagramPacket requestPacket = new DatagramPacket(new byte[64], 64);
                    serverSocket.receive(requestPacket);
                    System.out.println("receive........." + requestPacket.getPort());

                    final Integer packetID = Integer.valueOf(extractPacketID(requestPacket));
                    final int packetTypeID = getPacketTypeID(requestPacket);
                    // TM_ADDRESS_REQUEST_ID
                    // JM_LOOKUP_REQUEST_ID
                    System.out.println(packetTypeID + "...." + requestPacket.getSocketAddress());
                    if (packetTypeID == JM_LOOKUP_REQUEST_ID) {

                        System.out.println("Received job manager lookup request from " + requestPacket.getSocketAddress());
                        final DatagramPacket responsePacket = createJobManagerLookupReplyPacket(ipcPort,
                                magicNumber);
                        responsePacket.setAddress(requestPacket.getAddress());
                        responsePacket.setPort(requestPacket.getPort());

                        serverSocket.send(responsePacket);

                    } else if (packetTypeID == TM_ADDRESS_REQUEST_ID) {
                        System.out.println("Received task manager address request from " + requestPacket.getSocketAddress());
                        final DatagramPacket responsePacket = createTaskManagerAddressReplyPacket(requestPacket
                                .getAddress(), magicNumber);
                        responsePacket.setAddress(requestPacket.getAddress());
                        responsePacket.setPort(requestPacket.getPort());

                        serverSocket.send(responsePacket);

                    } else {
                        System.out.println("Received packet of unknown type " + packetTypeID + ", discarding...");
                    }


                }

            } catch (Exception e) {

            }
        }).start();




        System.in.read();
    }

    private static DatagramPacket createJobManagerLookupReplyPacket(final int ipcPort, final int magicNumber) {

        final byte[] bytes = new byte[16];
        integerToByteArray(magicNumber, MAGIC_NUMBER_OFFSET, bytes);
        integerToByteArray(generateRandomPacketID(), PACKET_ID_OFFSET, bytes);
        integerToByteArray(JM_LOOKUP_REPLY_ID, PACKET_TYPE_ID_OFFSET, bytes);
        integerToByteArray(ipcPort, PAYLOAD_OFFSET, bytes);

        return new DatagramPacket(bytes, bytes.length);
    }
    private static DatagramPacket createTaskManagerAddressRequestPacket(final int magicNumber) {

        final byte[] bytes = new byte[12];
        integerToByteArray(magicNumber, MAGIC_NUMBER_OFFSET, bytes);
        integerToByteArray(generateRandomPacketID(), PACKET_ID_OFFSET, bytes);
        integerToByteArray(TM_ADDRESS_REQUEST_ID, PACKET_TYPE_ID_OFFSET, bytes);

        return new DatagramPacket(bytes, bytes.length);
    }
    private static DatagramPacket createTaskManagerAddressReplyPacket(final InetAddress taskManagerAddress,
                                                                      final int magicNumber) {
        final byte[] addr = taskManagerAddress.getAddress();
        final byte[] bytes = new byte[20 + addr.length];
        integerToByteArray(magicNumber, MAGIC_NUMBER_OFFSET, bytes);
        integerToByteArray(generateRandomPacketID(), PACKET_ID_OFFSET, bytes);
        integerToByteArray(TM_ADDRESS_REPLY_ID, PACKET_TYPE_ID_OFFSET, bytes);
        integerToByteArray(addr.length, PAYLOAD_OFFSET, bytes);
        System.arraycopy(addr, 0, bytes, PAYLOAD_OFFSET + 4, addr.length);
        return new DatagramPacket(bytes, bytes.length);
    }



    private static int getPacketTypeID(final DatagramPacket packet) {
        final byte[] data = packet.getData();
        if (data == null) {
            return -1;
        }
        if (packet.getLength() < (PACKET_TYPE_ID_OFFSET + 4)) {
            return -1;
        }
        return byteArrayToInteger(data, PACKET_TYPE_ID_OFFSET);
    }

    private static int extractPacketID(final DatagramPacket packet) {
        final byte[] data = packet.getData();
        if (data == null) {
            return -1;
        }
        if (data.length < (PACKET_ID_OFFSET + 4)) {
            return -1;
        }
        return byteArrayToInteger(data, PACKET_ID_OFFSET);
    }


    private static final int MAGIC_NUMBER_OFFSET = 0;
    private static final int PACKET_ID_OFFSET = 4;
    private static final int PACKET_TYPE_ID_OFFSET = 8;
    private static final int PAYLOAD_OFFSET = 12;

    private static final int JM_LOOKUP_REQUEST_ID = 0;
    private static final int JM_LOOKUP_REPLY_ID = 1;
    private static final int TM_ADDRESS_REQUEST_ID = 2;
    private static final int TM_ADDRESS_REPLY_ID = 3;




    private static int generateRandomPacketID() {

        return (int) (Math.random() * (double) Integer.MAX_VALUE);
    }

    private static DatagramPacket createJobManagerLookupRequestPacket(final int magicNumber) {

        final byte[] bytes = new byte[12];
        integerToByteArray(magicNumber, MAGIC_NUMBER_OFFSET, bytes);
        integerToByteArray(generateRandomPacketID(), PACKET_ID_OFFSET, bytes);
        integerToByteArray(JM_LOOKUP_REQUEST_ID, PACKET_TYPE_ID_OFFSET, bytes);

        return new DatagramPacket(bytes, bytes.length);
    }


    private static void integerToByteArray(final int integerToSerialize, final int offset, final byte[] byteArray) {

        for (int i = 0; i < 4; ++i) {
            final int shift = i << 3; // i * 8
            byteArray[(offset + 3) - i] = (byte) ((integerToSerialize & (0xff << shift)) >>> shift);
        }
    }

    private static int byteArrayToInteger(final byte[] byteArray, final int offset) {

        int integer = 0;

        for (int i = 0; i < 4; ++i) {
            integer |= (byteArray[(offset + 3) - i] & 0xff) << (i << 3);
        }

        return integer;
    }

}
